#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
ROI设置对话框
"""

from PyQt5.QtWidgets import (
    QDialog, QVBoxLayout, QHBoxLayout, QPushButton, QLabel,
    QFrame, QSlider, QCheckBox
)
from PyQt5.QtCore import Qt, QRectF, pyqtSignal
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QColor

import cv2
import numpy as np

class ROIPreview(QFrame):
    """ROI预览窗口"""
    
    roi_changed = pyqtSignal(QRectF)
    
    def __init__(self, parent=None):
        """初始化"""
        super().__init__(parent)
        
        # 设置最小尺寸
        self.setMinimumSize(400, 300)
        
        # 设置样式
        self.setFrameStyle(QFrame.StyledPanel | QFrame.Sunken)
        self.setLineWidth(2)
        
        # 图像数据
        self.image = None
        self.pixmap = None
        
        # ROI数据
        self.roi = QRectF(0.25, 0.25, 0.5, 0.5)  # 默认是中心区域，相对坐标 (x, y, width, height)
        
        # 交互状态
        self.dragging = False
        self.drag_start = None
        self.drag_corner = None
        self.current_cursor = Qt.ArrowCursor
        
        # 追踪鼠标
        self.setMouseTracking(True)
    
    def set_image(self, image):
        """
        设置预览图像
        
        Args:
            image: OpenCV格式图像（BGR）
        """
        if image is None:
            return
        
        # 保存图像
        self.image = image
        
        # 转换为QPixmap
        height, width, channels = self.image.shape
        rgb_image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
        bytes_per_line = channels * width
        q_image = QImage(
            rgb_image.data, 
            width, 
            height, 
            bytes_per_line, 
            QImage.Format_RGB888
        )
        self.pixmap = QPixmap.fromImage(q_image)
        
        # 更新显示
        self.update()
    
    def set_roi(self, roi_rect):
        """
        设置ROI区域
        
        Args:
            roi_rect (QRectF): ROI矩形区域，相对坐标(0-1)
        """
        self.roi = roi_rect
        self.update()
        self.roi_changed.emit(self.roi)
    
    def paintEvent(self, event):
        """绘制事件"""
        super().paintEvent(event)
        
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        
        # 绘制图像
        if self.pixmap:
            # 计算缩放后图像的尺寸，保持纵横比
            frame_size = self.size()
            image_size = self.pixmap.size()
            
            scaled_size = image_size.scaled(
                frame_size, 
                Qt.KeepAspectRatio
            )
            
            # 计算居中位置
            left = (frame_size.width() - scaled_size.width()) / 2
            top = (frame_size.height() - scaled_size.height()) / 2
            
            # 绘制图像
            painter.drawPixmap(
                int(left), int(top),
                int(scaled_size.width()), int(scaled_size.height()),
                self.pixmap
            )
            
            # 计算ROI的实际像素坐标
            roi_rect = QRectF(
                left + self.roi.x() * scaled_size.width(),
                top + self.roi.y() * scaled_size.height(),
                self.roi.width() * scaled_size.width(),
                self.roi.height() * scaled_size.height()
            )
            
            # 绘制ROI区域
            painter.setPen(QPen(QColor(255, 0, 0), 2))
            painter.drawRect(roi_rect)
            
            # 绘制调整点
            corner_size = 8
            painter.setBrush(QColor(255, 0, 0))
            
            # 四个角的点
            # 左上角
            painter.drawRect(int(roi_rect.left() - corner_size/2), 
                             int(roi_rect.top() - corner_size/2), 
                             corner_size, corner_size)
            
            # 右上角
            painter.drawRect(int(roi_rect.right() - corner_size/2), 
                             int(roi_rect.top() - corner_size/2), 
                             corner_size, corner_size)
            
            # 左下角
            painter.drawRect(int(roi_rect.left() - corner_size/2), 
                             int(roi_rect.bottom() - corner_size/2), 
                             corner_size, corner_size)
            
            # 右下角
            painter.drawRect(int(roi_rect.right() - corner_size/2), 
                             int(roi_rect.bottom() - corner_size/2), 
                             corner_size, corner_size)
                             
            # 添加左右边缘的调整点
            # 左边中点
            painter.drawRect(int(roi_rect.left() - corner_size/2), 
                             int(roi_rect.top() + roi_rect.height()/2 - corner_size/2), 
                             corner_size, corner_size)
                             
            # 右边中点
            painter.drawRect(int(roi_rect.right() - corner_size/2), 
                             int(roi_rect.top() + roi_rect.height()/2 - corner_size/2), 
                             corner_size, corner_size)
                             
            # 上边中点
            painter.drawRect(int(roi_rect.left() + roi_rect.width()/2 - corner_size/2),
                             int(roi_rect.top() - corner_size/2),
                             corner_size, corner_size)
                             
            # 下边中点
            painter.drawRect(int(roi_rect.left() + roi_rect.width()/2 - corner_size/2),
                             int(roi_rect.bottom() - corner_size/2),
                             corner_size, corner_size)
                             
    def mousePressEvent(self, event):
        """鼠标按下事件"""
        if self.pixmap and event.button() == Qt.LeftButton:
            # 图像区域
            frame_size = self.size()
            image_size = self.pixmap.size()
            
            scaled_size = image_size.scaled(
                frame_size, 
                Qt.KeepAspectRatio
            )
            
            left = (frame_size.width() - scaled_size.width()) / 2
            top = (frame_size.height() - scaled_size.height()) / 2
            
            # 图像内有效区域
            image_rect = QRectF(
                left, top,
                scaled_size.width(), scaled_size.height()
            )
            
            # 点击位置
            pos = event.pos()
            
            # 检查是否点击在图像内
            if image_rect.contains(pos):
                # ROI的实际像素坐标
                roi_rect = QRectF(
                    left + self.roi.x() * scaled_size.width(),
                    top + self.roi.y() * scaled_size.height(),
                    self.roi.width() * scaled_size.width(),
                    self.roi.height() * scaled_size.height()
                )
                
                # 检查是否点击在ROI角落或边缘
                corner_size = 8
                
                # 左上角
                lt_rect = QRectF(
                    roi_rect.left() - corner_size/2,
                    roi_rect.top() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 右上角
                rt_rect = QRectF(
                    roi_rect.right() - corner_size/2,
                    roi_rect.top() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 左下角
                lb_rect = QRectF(
                    roi_rect.left() - corner_size/2,
                    roi_rect.bottom() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 右下角
                rb_rect = QRectF(
                    roi_rect.right() - corner_size/2,
                    roi_rect.bottom() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 左边中点
                l_rect = QRectF(
                    roi_rect.left() - corner_size/2,
                    roi_rect.top() + roi_rect.height()/2 - corner_size/2,
                    corner_size, corner_size
                )
                
                # 右边中点
                r_rect = QRectF(
                    roi_rect.right() - corner_size/2,
                    roi_rect.top() + roi_rect.height()/2 - corner_size/2,
                    corner_size, corner_size
                )
                
                # 上边中点
                t_rect = QRectF(
                    roi_rect.left() + roi_rect.width()/2 - corner_size/2,
                    roi_rect.top() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 下边中点
                b_rect = QRectF(
                    roi_rect.left() + roi_rect.width()/2 - corner_size/2,
                    roi_rect.bottom() - corner_size/2,
                    corner_size, corner_size
                )
                
                if lt_rect.contains(pos):
                    self.dragging = True
                    self.drag_corner = "lt"
                    self.drag_start = pos
                elif rt_rect.contains(pos):
                    self.dragging = True
                    self.drag_corner = "rt"
                    self.drag_start = pos
                elif lb_rect.contains(pos):
                    self.dragging = True
                    self.drag_corner = "lb"
                    self.drag_start = pos
                elif rb_rect.contains(pos):
                    self.dragging = True
                    self.drag_corner = "rb"
                    self.drag_start = pos
                elif l_rect.contains(pos):
                    self.dragging = True
                    self.drag_corner = "l"
                    self.drag_start = pos
                elif r_rect.contains(pos):
                    self.dragging = True
                    self.drag_corner = "r"
                    self.drag_start = pos
                elif t_rect.contains(pos):
                    self.dragging = True
                    self.drag_corner = "t"
                    self.drag_start = pos
                elif b_rect.contains(pos):
                    self.dragging = True
                    self.drag_corner = "b"
                    self.drag_start = pos
                elif roi_rect.contains(pos):
                    # 拖动整个ROI
                    self.dragging = True
                    self.drag_corner = "move"
                    self.drag_start = pos
                else:
                    # 创建新ROI
                    rel_x = (pos.x() - left) / scaled_size.width()
                    rel_y = (pos.y() - top) / scaled_size.height()
                    
                    # 从点击位置开始创建ROI
                    self.roi = QRectF(
                        max(0, min(rel_x, 1.0)),
                        max(0, min(rel_y, 1.0)),
                        0.0, 0.0
                    )
                    
                    self.dragging = True
                    self.drag_corner = "rb"
                    self.drag_start = pos
    
    def mouseMoveEvent(self, event):
        """鼠标移动事件"""
        if self.pixmap:
            # 图像区域
            frame_size = self.size()
            image_size = self.pixmap.size()
            
            scaled_size = image_size.scaled(
                frame_size, 
                Qt.KeepAspectRatio
            )
            
            left = (frame_size.width() - scaled_size.width()) / 2
            top = (frame_size.height() - scaled_size.height()) / 2
            
            # 图像内有效区域
            image_rect = QRectF(
                left, top,
                scaled_size.width(), scaled_size.height()
            )
            
            # 鼠标位置
            pos = event.pos()
            
            # 查找鼠标悬停在ROI的哪个位置以更新光标
            if not self.dragging and image_rect.contains(pos):
                # ROI的实际像素坐标
                roi_rect = QRectF(
                    left + self.roi.x() * scaled_size.width(),
                    top + self.roi.y() * scaled_size.height(),
                    self.roi.width() * scaled_size.width(),
                    self.roi.height() * scaled_size.height()
                )
                
                # 检查鼠标是否在某个调整点上
                corner_size = 8
                
                # 左上角
                lt_rect = QRectF(
                    roi_rect.left() - corner_size/2,
                    roi_rect.top() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 右上角
                rt_rect = QRectF(
                    roi_rect.right() - corner_size/2,
                    roi_rect.top() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 左下角
                lb_rect = QRectF(
                    roi_rect.left() - corner_size/2,
                    roi_rect.bottom() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 右下角
                rb_rect = QRectF(
                    roi_rect.right() - corner_size/2,
                    roi_rect.bottom() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 左边中点
                l_rect = QRectF(
                    roi_rect.left() - corner_size/2,
                    roi_rect.top() + roi_rect.height()/2 - corner_size/2,
                    corner_size, corner_size
                )
                
                # 右边中点
                r_rect = QRectF(
                    roi_rect.right() - corner_size/2,
                    roi_rect.top() + roi_rect.height()/2 - corner_size/2,
                    corner_size, corner_size
                )
                
                # 上边中点
                t_rect = QRectF(
                    roi_rect.left() + roi_rect.width()/2 - corner_size/2,
                    roi_rect.top() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 下边中点
                b_rect = QRectF(
                    roi_rect.left() + roi_rect.width()/2 - corner_size/2,
                    roi_rect.bottom() - corner_size/2,
                    corner_size, corner_size
                )
                
                # 根据鼠标位置设置光标形状
                if lt_rect.contains(pos) or rb_rect.contains(pos):
                    self.setCursor(Qt.SizeFDiagCursor)  # 左上-右下调整
                    self.current_cursor = Qt.SizeFDiagCursor
                elif rt_rect.contains(pos) or lb_rect.contains(pos):
                    self.setCursor(Qt.SizeBDiagCursor)  # 右上-左下调整
                    self.current_cursor = Qt.SizeBDiagCursor
                elif l_rect.contains(pos) or r_rect.contains(pos):
                    self.setCursor(Qt.SizeHorCursor)  # 水平调整
                    self.current_cursor = Qt.SizeHorCursor
                elif t_rect.contains(pos) or b_rect.contains(pos):
                    self.setCursor(Qt.SizeVerCursor)  # 垂直调整
                    self.current_cursor = Qt.SizeVerCursor
                elif roi_rect.contains(pos):
                    self.setCursor(Qt.SizeAllCursor)  # 移动
                    self.current_cursor = Qt.SizeAllCursor
                else:
                    self.setCursor(Qt.ArrowCursor)  # 默认
                    self.current_cursor = Qt.ArrowCursor
            
            # 处理拖动
            if self.dragging and self.drag_start:
                # 设置鼠标样式
                if self.drag_corner in ["lt", "rb"]:
                    self.setCursor(Qt.SizeFDiagCursor)
                elif self.drag_corner in ["rt", "lb"]:
                    self.setCursor(Qt.SizeBDiagCursor)
                elif self.drag_corner in ["l", "r"]:
                    self.setCursor(Qt.SizeHorCursor)
                elif self.drag_corner in ["t", "b"]:
                    self.setCursor(Qt.SizeVerCursor)
                elif self.drag_corner == "move":
                    self.setCursor(Qt.SizeAllCursor)
                
                # 计算鼠标移动的相对坐标变化
                dx = (pos.x() - self.drag_start.x()) / scaled_size.width()
                dy = (pos.y() - self.drag_start.y()) / scaled_size.height()
                
                # 当前ROI
                roi = QRectF(self.roi)
                
                # 根据拖动类型调整ROI
                if self.drag_corner == "lt":  # 左上角
                    new_x = roi.x() + dx
                    new_y = roi.y() + dy
                    new_w = roi.width() - dx
                    new_h = roi.height() - dy
                    
                    # 确保宽高不为负
                    if new_w > 0 and new_h > 0:
                        roi.setX(new_x)
                        roi.setY(new_y)
                        roi.setWidth(new_w)
                        roi.setHeight(new_h)
                        
                elif self.drag_corner == "rt":  # 右上角
                    new_y = roi.y() + dy
                    new_w = roi.width() + dx
                    new_h = roi.height() - dy
                    
                    # 确保宽高不为负
                    if new_w > 0 and new_h > 0:
                        roi.setY(new_y)
                        roi.setWidth(new_w)
                        roi.setHeight(new_h)
                        
                elif self.drag_corner == "lb":  # 左下角
                    new_x = roi.x() + dx
                    new_w = roi.width() - dx
                    new_h = roi.height() + dy
                    
                    # 确保宽高不为负
                    if new_w > 0 and new_h > 0:
                        roi.setX(new_x)
                        roi.setWidth(new_w)
                        roi.setHeight(new_h)
                        
                elif self.drag_corner == "rb":  # 右下角
                    new_w = roi.width() + dx
                    new_h = roi.height() + dy
                    
                    # 确保宽高不为负
                    if new_w > 0 and new_h > 0:
                        roi.setWidth(new_w)
                        roi.setHeight(new_h)
                
                elif self.drag_corner == "l":  # 左边
                    new_x = roi.x() + dx
                    new_w = roi.width() - dx
                    
                    # 确保宽度不为负
                    if new_w > 0:
                        roi.setX(new_x)
                        roi.setWidth(new_w)
                        
                elif self.drag_corner == "r":  # 右边
                    new_w = roi.width() + dx
                    
                    # 确保宽度不为负
                    if new_w > 0:
                        roi.setWidth(new_w)
                        
                elif self.drag_corner == "t":  # 上边
                    new_y = roi.y() + dy
                    new_h = roi.height() - dy
                    
                    # 确保高度不为负
                    if new_h > 0:
                        roi.setY(new_y)
                        roi.setHeight(new_h)
                        
                elif self.drag_corner == "b":  # 下边
                    new_h = roi.height() + dy
                    
                    # 确保高度不为负
                    if new_h > 0:
                        roi.setHeight(new_h)
                        
                elif self.drag_corner == "move":  # 整体移动
                    new_x = roi.x() + dx
                    new_y = roi.y() + dy
                    
                    roi.moveLeft(new_x)
                    roi.moveTop(new_y)
                
                # 确保ROI在图像范围内
                if roi.left() < 0:
                    roi.moveLeft(0)
                if roi.top() < 0:
                    roi.moveTop(0)
                if roi.right() > 1:
                    roi.moveRight(1)
                if roi.bottom() > 1:
                    roi.moveBottom(1)
                
                # 更新ROI
                self.roi = roi
                
                # 更新起始拖动点
                self.drag_start = pos
                
                # 更新显示
                self.update()
                
                # 发射ROI变更信号
                self.roi_changed.emit(self.roi)
    
    def mouseReleaseEvent(self, event):
        """鼠标释放事件"""
        if event.button() == Qt.LeftButton and self.dragging:
            self.dragging = False
            self.drag_start = None
            self.drag_corner = None
            
            # 确保ROI不为空
            if self.roi.width() < 0.05 or self.roi.height() < 0.05:
                # 恢复默认ROI
                self.roi = QRectF(0.25, 0.25, 0.5, 0.5)
                self.update()
                self.roi_changed.emit(self.roi)

class ROIDialog(QDialog):
    """ROI设置对话框"""
    
    def __init__(self, image, current_roi=None, parent=None):
        """
        初始化对话框
        
        Args:
            image: OpenCV格式图像
            current_roi (QRectF, optional): 当前ROI设置，相对坐标(0-1)
            parent (QWidget, optional): 父窗口
        """
        super().__init__(parent)
        
        # 保存图像和ROI
        self.image = image
        self.roi = current_roi or QRectF(0.25, 0.25, 0.5, 0.5)
        
        # 设置对话框属性
        self.setWindowTitle("设置ROI区域")
        self.resize(600, 500)
        
        # 初始化UI
        self._setup_ui()
    
    def _setup_ui(self):
        """设置界面"""
        # 主布局
        main_layout = QVBoxLayout(self)
        
        # 顶部说明文字
        main_layout.addWidget(QLabel("拖动矩形框或角点来设置ROI区域。此区域将应用于所有图像。"))
        
        # 预览窗口
        self.preview = ROIPreview()
        self.preview.set_image(self.image)
        self.preview.set_roi(self.roi)
        self.preview.roi_changed.connect(self._on_roi_changed)
        main_layout.addWidget(self.preview)
        
        # 应用于所有图像的复选框
        self.apply_all_checkbox = QCheckBox("应用于数据集中的所有图像")
        self.apply_all_checkbox.setChecked(True)
        main_layout.addWidget(self.apply_all_checkbox)
        
        # 按钮布局
        button_layout = QHBoxLayout()
        
        # 重置按钮
        self.reset_button = QPushButton("重置")
        self.reset_button.clicked.connect(self._on_reset)
        button_layout.addWidget(self.reset_button)
        
        # 取消按钮
        self.cancel_button = QPushButton("取消")
        self.cancel_button.clicked.connect(self.reject)
        button_layout.addWidget(self.cancel_button)
        
        # 确定按钮
        self.ok_button = QPushButton("确定")
        self.ok_button.clicked.connect(self.accept)
        self.ok_button.setDefault(True)
        button_layout.addWidget(self.ok_button)
        
        main_layout.addLayout(button_layout)
    
    def _on_roi_changed(self, roi):
        """ROI变更事件"""
        self.roi = roi
    
    def _on_reset(self):
        """重置ROI"""
        self.roi = QRectF(0.25, 0.25, 0.5, 0.5)
        self.preview.set_roi(self.roi)
    
    def get_roi(self):
        """
        获取用户设置的ROI
        
        Returns:
            tuple: (QRectF, bool) - ROI和是否应用于所有图像
        """
        return (self.roi, self.apply_all_checkbox.isChecked()) 